# from rocketmq.client import Producer, Message
from osgeo import gdal,gdalconst
from xml.dom.minidom import parse
import datetime as dt
import json
import daemon
import redis
import psycopg2
import shutil
import zipfile
import time
import os
import re


def table_2_data(cargs,name,tif_path):
    connection = psycopg2.connect(**cargs)
    connection.autocommit = True
    cursor = connection.cursor()
    query = f"""select tile,product_date,cloud,ST_AsText(geom)
        from sentinel2_l1c
        where name = '{name}';
        """
    cursor.execute(query)
    data = cursor.fetchone()
    cursor.close()
    connection.close()
    result = {
        "tile": data[0],  # 假设其他属性字段在结果中的索引为1
        "product_date": data[1].strftime("%Y-%m-%d"),
        "cpjbdm": "1",
        "cpjdmc": "L1C",
        "imgtype_code": "S2",
        "imgtype_name": "Sentinel-2",
        "cloud": float(data[2]),
        "tif_path": tif_path,
        "geometry": data[3]}
    return result


def send_message(json_dir,zip_path,data,redisr,topic="henan_rsanalysis_s2_1c_collect",addr='192.168.1.70:9876'):
    name = os.path.basename(zip_path).replace(".zip","")
    json_path = os.path.join(json_dir,f"{name}.json")
    with open(json_path,"w") as json_file:
        json.dump(data,json_file,indent=4)
    redisr.rpush('raster_json_download', f'"{json_path}"')
    # producer = Producer("data")
    # producer.set_name_server_address(addr)
    # producer.start()
    # message_content = {"file_path":json_path}
    # message_body = json.dumps(message_content)
    # msg = Message(topic)
    # msg.set_tags("TagA")
    # msg.set_body(message_body)
    # _ = producer.send_sync(msg)
    # producer.shutdown()
    return 


def zip_2_jp2(zip_path, bands, save_dir):
    band_names = {
        "1":"T.*_B01\.jp2",
        "2":"T.*_B02\.jp2",
        "3":"T.*_B03\.jp2",
        "4":"T.*_B04\.jp2",
        "5":"T.*_B05\.jp2",
        "6":"T.*_B06\.jp2",
        "7":"T.*_B07\.jp2",
        "8A":"T.*_B8A\.jp2",
        "8":"T.*_B08\.jp2",
        "9":"T.*_B09\.jp2",
        "10":"T.*_B10\.jp2",
        "11":"T.*_B11\.jp2",
        "12":"T.*_B12\.jp2"}
    pathres = list(map(lambda x:".*/IMG_DATA/"+band_names[f"{x}"],bands))
    paths = []
    zipFile = zipfile.ZipFile(zip_path, "r")
    for pathre in pathres:
        for info in zipFile.infolist():
            if re.match(pathre,info.filename):
                path = zipFile.extract(info.filename, save_dir)
                paths.append(path)
    
    zipFile.close()
    return paths
    

def jp2s_2_tif(jp2_paths,tif_path,nodata=0):
    vrt_path = tif_path.replace('.tif','.vrt')
    options = gdal.BuildVRTOptions(
        resolution="highest",
        VRTNodata=nodata,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        separate=True)
    gdal.BuildVRT(vrt_path,jp2_paths,options=options)
    creationOptions = [
        "BIGTIFF=YES",
        "TILED=YES",
        "BLOCKXSIZE=1024",
        # "TFW=YES",
        "NUM_THREADS=ALL_CPUS"]
    options = gdal.TranslateOptions(
        format="GTiff",
        xRes=10,
        yRes=10,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        creationOptions=creationOptions)
    gdal.Translate(tif_path,vrt_path,options=options)
    os.remove(vrt_path)
    return


def zip_2_tif(zip_path,save_dir,bands=['2','3','4','5','6','7',"8A",'8','11','12']):
    zip_name = os.path.basename(zip_path)
    year_month = zip_name.split("_")[2][:6]
    date_dir = os.path.join(save_dir,year_month)
    if not os.path.exists(date_dir):os.makedirs(date_dir)
    tif_name = zip_name.replace('.SAFE', '').replace('.zip', '.tif')
    tif_path = os.path.join(date_dir, tif_name)
    # if os.path.exists(tif_path): return tif_path
    jp2_paths = zip_2_jp2(zip_path,bands,date_dir)
    jp2s_2_tif(jp2_paths,tif_path)
    unzip_dir = tif_path.replace(".tif",".SAFE")
    if os.path.exists(unzip_dir): shutil.rmtree(unzip_dir)
    return tif_path


def main():
    with open('/data/jiabing/RS_CODE/Sentinel2download/config.json',"r") as file:
        params = json.load(file)
    tif_dir = params["tif_dir"]
    json_dir = params['json_dir']
    cargs = params["cargs"]
    rediscon = params["rediscon"]
    pool = redis.ConnectionPool(**rediscon)
    redisr = redis.Redis(connection_pool=pool)
    while True:
        time.sleep(10)
        _,element = redisr.blpop('downloader_raster',0)
        zip_path = element.decode('utf-8').replace('"','')
        print(os.path.basename(zip_path).replace(".zip",""))
        tif_path = zip_2_tif(zip_path, tif_dir)
        name = os.path.basename(tif_path).replace(".tif","")
        data = table_2_data(cargs,name,tif_path)
        send_message(json_dir,zip_path,data,redisr)
    return
    

if __name__ == "__main__":
    # with open('/data/Yang/project_1/config.json',"r") as file:
    #     params = json.load(file)
    # cargs = params["cargs"]
    # name = r"S2A_MSIL1C_20240401T030521_N0510_R075_T49SER_20240401T054712"
    # tif_path = r"S2A_MSIL1C_20240401T030521_N0510_R075_T49SER_20240401T054712.tif"
    # table_2_data(cargs,name,tif_path)
    # main()
    with daemon.DaemonContext(
        stdout=open(r"/data/logfile/RasterDaemon_out.log","w"),
        stderr=open(r"/data/logfile/RasterDaemon_err.log","w")):
        main()